# coding: UTF-8
import json
import time
import xml.dom.minidom


class TemplateXmlAndJsonWriter:
    def __init__(
            self,
            image_info,
            databaseSrc='Unknown',
            localImgPath=None,
            shape_type=None):
        self.filename = image_info['image_name']
        self.databaseSrc = databaseSrc
        self.image_info = image_info
        self.boxlist = []
        self.dict_list=[]
        self.localImgPath = localImgPath
        self.shape_type = shape_type

    def addPolygon(self, bpoint, label, annotator, instance_id):
        polygon = {}

        polygon['polygon'] = bpoint
        polygon['label'] = label
        polygon['point_num'] = str(len(bpoint))
        polygon['id'] = instance_id
        polygon['deleted'] = 0
        polygon['draw'] = True
        polygon['shape_type'] = self.shape_type
        polygon['verified'] = 0
        polygon['user'] = annotator
        date = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(time.time()))
        polygon['date'] = date
        self.dict_list.append(polygon)

    def save_json(self, targetFile=None):

        djson = {"filename": self.filename,
                 "imgHeight": self.image_info['height'],
                 "imgWidth": self.image_info['width'],
                 "depth": self.image_info['depth'],
                 "objects": self.dict_list
                 }

        with open(targetFile, "w") as f:
            json.dump(djson, f, sort_keys=True, indent=4)
        print("%s saved."%targetFile)
    
    def save_xml(self, results, xmlPath):

        impl = xml.dom.minidom.getDOMImplementation()
        dom = impl.createDocument(None, 'annotation', None)
        root = dom.documentElement

        annotator = dom.createElement('annotator')
        root.appendChild(annotator)
        annotator_name = dom.createTextNode(self.image_info['annotator'])
        annotator.appendChild(annotator_name)

        folder = dom.createElement('folder')
        root.appendChild(folder)
        name_folfer = dom.createTextNode('Unknown')
        folder.appendChild(name_folfer)

        filename = dom.createElement('filename')
        root.appendChild(filename)
        name_img = dom.createTextNode(self.image_info['image_name'])
        filename.appendChild(name_img)

        filepath = dom.createElement('path')
        root.appendChild(filepath)
        path_img = dom.createTextNode(self.image_info['image_name'])
        filepath.appendChild(path_img)

        source = dom.createElement('source')
        root.appendChild(source)
        database = dom.createElement('database')
        database_name = dom.createTextNode('Unknown')
        database.appendChild(database_name)
        source.appendChild(database)

        img_size = dom.createElement('size')
        root.appendChild(img_size)
        width = dom.createElement('width')
        width_num = dom.createTextNode(str(self.image_info['width']))
        width.appendChild(width_num)
        height = dom.createElement('height')
        height_num = dom.createTextNode(str(self.image_info['height']))
        height.appendChild(height_num)
        depth = dom.createElement('depth')
        depth_num = dom.createTextNode(str(self.image_info['depth']))
        depth.appendChild(depth_num)
        img_size.appendChild(width)
        img_size.appendChild(height)
        img_size.appendChild(depth)

        segmented = dom.createElement('segmented')
        root.appendChild(segmented)
        segmented_num = dom.createTextNode('0')
        segmented.appendChild(segmented_num)

        for i in range(len(results)):
            img_object = dom.createElement('object')
            root.appendChild(img_object)
            label_name = dom.createElement('name')
            namecls = dom.createTextNode(results[i]['class'])
            label_name.appendChild(namecls)
            pose = dom.createElement('pose')
            pose_name = dom.createTextNode('Unspecified')
            pose.appendChild(pose_name)
            truncated = dom.createElement('truncated')
            truncated_num = dom.createTextNode('0')
            truncated.appendChild(truncated_num)
            difficult = dom.createElement('difficult')
            difficult_num = dom.createTextNode('0')
            difficult.appendChild(difficult_num)
            bndbox = dom.createElement('bndbox')
            xmin = dom.createElement('xmin')
            xmin_num = dom.createTextNode(str(int(results[i]['bbox'][0])))
            xmin.appendChild(xmin_num)
            ymin = dom.createElement('ymin')
            ymin_num = dom.createTextNode(str(int(results[i]['bbox'][1])))
            ymin.appendChild(ymin_num)
            xmax = dom.createElement('xmax')
            xmax_num = dom.createTextNode(str(int(results[i]['bbox'][2])))
            xmax.appendChild(xmax_num)
            ymax = dom.createElement('ymax')
            ymax_num = dom.createTextNode(str(int(results[i]['bbox'][3])))
            ymax.appendChild(ymax_num)
            bndbox.appendChild(xmin)
            bndbox.appendChild(ymin)
            bndbox.appendChild(xmax)
            bndbox.appendChild(ymax)
            img_object.appendChild(label_name)
            img_object.appendChild(pose)
            img_object.appendChild(truncated)
            img_object.appendChild(difficult)
            img_object.appendChild(bndbox)

        f = open(xmlPath, 'w')
        dom.writexml(f, addindent='  ', newl='\n')
        f.close()
        print("%s saved." % xmlPath)